[Experimental] MiniMax-H3 LoRA training on B200: 2.5x faster step from three changes, with takeaways for the repo's acceleration - #1680

Setup and Key Improvement

Hardware. One node, 8 × NVIDIA B200 (sm_100, 183 GB each), driver 580.126.20, CUDA 13.0. Weights on node-local NVMe.

Env installation. torch 2.14.0+cu130 (cuDNN bundled), flash-attn 2.8.4 built from source for sm_100, peft 0.20.0, accelerate 1.14.0, DiffSynth-Studio main at ce9f454, editable install.

Script. examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh, stage 2 (--task sft:train). Flags as in the script except the step count; 8-process DDP through accelerate. The full command is under Reproduce.

Model and job. MiniMax-H3 FL2VA DiT (50 blocks, about 33 B parameters, bf16, frozen) with LoRA r32 on the attention and MLP projections. Each step trains on 124 frames at 480×832 with audio, a packed sequence of 14 912 tokens, with gradient checkpointing on.

Measurement. 30 optimizer steps per process (--dataset_repeat 240); seeds fixed to 42 + rank by a launcher wrapper, since train.py does not seed. Step time is the median of steps 6 to 30 from CUDA events on rank 0; all ranks agree within 1 ms. Profiles cover steps 8 to 10 on all 8 ranks. The unmodified step takes 6.32 s with the GPU 99% busy.

Improvements in short. Three changes to diffsynth/models/minimax_h3_dit.py, training step 6.32 s → 2.53 s (2.5×).

What this PR does

Three commits, one file (diffsynth/models/minimax_h3_dit.py), limited to the H3 model so nothing else changes.

Peak memory per rank goes from 75.4 to 83.6 GB. The first step takes about 30 s with a cold inductor cache and about 12 s with a warm one; it was 7 s before. Model, precision (mixed_precision: 'no', bf16), world size and LoRA configuration are unchanged.
Kernel time per step per rank went from attention 3.85 s, element-wise 1.29 s, GEMM 1.16 s to 0.99 s, 0.31 s, 1.26 s (GEMMs on cuBLAS on both sides); exposed all-reduce 40 → 37 ms, idle 54 → 15 ms.

Traces (torch.profiler, steps 8 to 10, all 8 ranks; open in https://ui.perfetto.dev):

• baseline: https://drive.google.com/open?id=1LKNGMZntNh7lB4rS9hYJfF26kG8qemyX
• optimized: https://drive.google.com/open?id=1cQvISI_4vtQr7F0pX0f-grpWPDHnJZ2z

General takeaway for the repo's acceleration

Five takeaways, each with its cause and a fix that does not depend on this PR.

1. The attention dispatch ignores the GPU. diffsynth/core/attention/attention.py:61-80 picks the implementation once at import, from which packages import, in a fixed order: custom kernel, FA4 (flash_attn.cute), FA3 (flash_attn_interface), FA2 (flash_attn), SageAttention, xFormers, torch SDPA. On sm_90+ the FA2 kernels are an sm80 design: about 440 TFLOP/s on B200 at this shape, against 1.3 to 1.5 PFLOP/s for torch SDPA's cuDNN backend with the same error. In this job that was 61% of every step. Fix: in initialize_attention_priority(), read torch.cuda.get_device_capability(); on major version 9 or higher with neither FA3 nor FA4 importable, return the torch path and run torch_sdpa under sdpa_kernel([CUDNN_ATTENTION, FLASH_ATTENTION, EFFICIENT_ATTENTION, MATH], set_priority=True). FA2 stays the choice on sm_80 and older.

2. The choice is invisible. ATTENTION_IMPLEMENTATION is a module constant that nothing prints; the only way to notice a bad pick is a profile. Fix: log the selected implementation and the GPU name once at import.

3. The docs steer users to FA2 and contradict the repo's own advice. docs/en/Model_Details/Wan.md:175 (pip install flash-attn --no-build-isolation) and docs/en/Pipeline_Usage/Accelerated_Inference.md:20 (pip install "xfuser[flash-attn]>=0.4.3") install FA2; PyPI flash-attn has no FA4 subpackage and no doc gives an FA3 or FA4 step. docs/en/API_Reference/core/attention.md recommends the native PyTorch path with no extra packages, and because every model shares the one dispatch, installing flash-attn for Wan multi-GPU inference silently overrides that for every other model. Fix: the Wan and multi-GPU pages should say what installing flash-attn does to the other models and which flash-attn fits which GPU, or point Hopper and Blackwell users to the native path.

4. Un-fused element-wise work in the DiT block is generic. RoPE, the AdaLN index_select gathers, RMSNorm, SwiGLU and the residual gates were about 600 memory-bound launches per block execution and 20% of the step; torch.compile of the block body fused them into about 12 kernels. Any DiT block with AdaLN modulation has the same structure. Fix: the same compile, decided per model.

5. Gradient checkpointing recomputes attention for nothing. With fast attention, the block recompute's most expensive op is the attention forward, whose output costs about 240 MB per block to keep. Selective activation checkpointing that saves only that output cut another 12% here for 8 GB per rank. Same applies to any checkpointed DiT block. Fix: a selective policy, decided per model against the memory budget.

I can do 1 and 2 in this PR or in a separate one; 4 and 5 are in this PR for H3 only.

Correctness verification

To make sure the optimizations do not change what the model computes, I instrumented the points listed below to record intermediate values, and required every optimization to keep those values within a tolerance. The tool is probe.

"Actual difference to baseline" is the largest difference between this branch and the unmodified code over all listed steps and all 8 ranks, relative unless marked absolute. "Tolerance" is the allowed difference, set to about 3× the spread of three runs of the unmodified code.

Reproduce

Model. MiniMax-H3 FL2VA DiT: 50 blocks, hidden 5376, 56 heads × 128, FFN 14336, about 33 B parameters, bf16, frozen. LoRA rank 32 on attn.qkv_proj, attn.out_proj, mlp.fc1, mlp.fc2 (416 tensors, 155 M parameters, bf16). AdamW, lr 1e-4, gradient checkpointing per DiT block, 8-process DDP. One sample per step per process: 124 frames at 480×832 with audio, packed into 14 912 tokens (14 430 video, 414 audio, 68 text). The loss is the flow-matching MSE on video and audio.

Stage 1 (once): --task sft:data_process from the same .sh; it caches the text-encoder and VAE outputs of the example clip.

Stage 2 (the measured job; every flag is the script's, only --dataset_repeat changed so that each process does 30 steps):

accelerate launch --config_file <accelerate config> --num_processes 8 examples/minimax_h3/model_training/train.py \
  --dataset_base_path <stage-1 cache> --data_file_keys "video,input_audio" --extra_inputs "input_audio" \
  --height 480 --width 832 --num_frames 124 --learning_rate 1e-4 --remove_prefix_in_ckpt "pipe.dit." \
  --lora_base_model dit --lora_target_modules "attn.qkv_proj,attn.out_proj,mlp.fc1,mlp.fc2" --lora_rank 32 \
  --use_gradient_checkpointing --dataset_repeat 240 --num_epochs 1 \
  --model_id_with_origin_paths "MiniMax/MiniMax-H3:FL2VA/transformer/model*.safetensors" \
  --output_path <out> --find_unused_parameters --enable_csv_log --task sft:train

accelerate config: compute_environment: LOCAL_MACHINE, distributed_type: MULTI_GPU, num_processes: 8, mixed_precision: 'no', rdzv_backend: static. The repo ships no config for the LoRA path; this mirrors its ZeRO-3 yaml with plain DDP. bf16 there would mean autocast under MULTI_GPU, which fails FlashAttention's dtype check.

Seeds. train.py does not seed. The runs used a one-line launcher that calls torch.manual_seed(42 + rank) (and seeds numpy and random) before running train.py unchanged, so LoRA init, timesteps and noise are reproducible.

Timing. Median of steps 6 to 30 from CUDA events around each optimizer step on rank 0. Steps 1 to 5 are warm-up: first cuBLAS, cuDNN and NCCL calls, plus the inductor compile on the optimized side. The tqdm progress bar gives the same numbers within 1 ms. Profiles: torch.profiler with CPU and CUDA activities, no stacks, shapes or memory, over steps 8 to 10, one chrome trace per rank (schedule: wait 6, warmup 1, active 3).